import os
import torch
import cv2
from one_layer_model import IlluNet_with_Quad
import pyiqa
import os

def to_deploy(mo):
    for m in mo.modules():
        if hasattr(m, 'switch_to_deploy'):
            m.switch_to_deploy()
    return mo


device = torch.device('cuda')
niqe_metric = pyiqa.create_metric('niqe', device=device).eval()
PI_metric = pyiqa.create_metric('pi', device=device).eval()

pre_weight = './weight/sclm_noref.pth'
img_p = './input/3.bmp'

model = IlluNet_with_Quad(3,3,None).eval()
model.load_state_dict(torch.load(pre_weight))

model = to_deploy(model).to(device)

img_np = cv2.imread(img_p).astype('float32')/255.0
img_ten = torch.from_numpy(img_np).permute(2,0,1).unsqueeze(0).to(device)

with torch.no_grad():
    out,_ = model(img_ten)
    enhanced_np = out.detach().cpu().squeeze(0).permute(1,2,0).numpy()*255.0
    
    cv2.imwrite('./out/3.bmp',enhanced_np)